import sys
import os
from glob import glob
import os
import sys
import pandas as pd
import time

root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(root)

data = os.path.join(root,"datasets")
import main_v2.tools as tools

def main():

    baidu = os.path.join(data,"baidu")
    google = os.path.join(data,"google")
    bing = os.path.join(data,"bing")
    grd = os.path.join(data,"grd_test")
    tid = os.path.join(data, "tid")

    baidu_images = glob(os.path.join(baidu,"*/*.jpeg"))+glob(os.path.join(baidu,"*/*.png"))
    google_images = glob(os.path.join(google, "*/*.jpeg")) + glob(os.path.join(google, "*/*.png"))
    bing_images = glob(os.path.join(bing, "*/*.jpeg")) + glob(os.path.join(bing, "*/*.png"))
    grd_images = glob(os.path.join(grd,"*/*.jpeg"))+glob(os.path.join(grd,"*/*.jpg"))
    tid_images = glob(os.path.join(tid, "*/*.bmp"))

    baidu_images = sorted(baidu_images)
    google_images = sorted(google_images)
    bing_images = sorted(bing_images)
    grd_images = sorted(grd_images)
    tid_images = sorted(tid_images)

    if not os.path.exists(os.path.join(root,"outputs")):
        os.makedirs(os.path.join(root,"outputs"))

    funcs = {
        "entropy":tools.entropy, # 熵函数
        "EAV":tools.EAV, # EAV点锐度算法函数
        "SMD":tools.SMD, # SMD（灰度方差）函数
        "SMD2":tools.SMD2, # SMD2 （灰度方差乘积）函数
        "variance":tools.variance, # 方差函数
        "tenengrad":tools.tenengrad, # Tenengrad 梯度函数
        "vollath":tools.vollath,
        "laplacian":tools.laplacian,
        "brenner":tools.brenner,
        "nrss":tools.nrss,
        "energy":tools.energy,
        "NR_IQA":tools.NR_IQA,
        "JPEQ" : tools.JPEQ,
        "JPEQ2":tools.JPEQ2
    }

    # tid_dict = {"images": []}
    # for image in tid_images:
    #     print(image)
    #     tid_dict["images"].append(os.path.split(image)[-1])
    #     try:
    #         for key, val in funcs.items():
    #             score = val(image)
    #             if key not in tid_dict:
    #                 tid_dict[key] = [score]
    #             else:
    #                 tid_dict[key].append(score)
    #     except Exception as e:
    #         print(e)
    # df = pd.DataFrame(tid_dict)
    # df.to_csv(os.path.join(root, "outputs", "tid_JPEQ2.csv"), index=False)

    grd_dict = {"images":[]}
    for image in grd_images:
        grd_dict["images"].append(os.path.split(image)[-1])
        try:
            for key,val in funcs.items():

                start = time.time()
                score = val(image)
                end = time.time()

                t = end-start

                print(image, key,score,t)
                if key not in grd_dict:
                    grd_dict[key] = [score]
                    grd_dict["{}_{}".format(key,"time")] = [t]
                else:
                    grd_dict[key].append(score)
                    grd_dict["{}_{}".format(key, "time")].append(t)
        except Exception as e:
            print(e)
    df = pd.DataFrame(grd_dict)
    df.to_csv(os.path.join(root,"outputs","grd_test_306_196.csv"),index=False)

    # baidu_dict = {"images":[]}
    # for image in baidu_images:
    #     print(image)
    #     baidu_dict["images"].append(os.path.split(image)[-1])
    #     try:
    #         for key,val in funcs.items():
    #             score = val(image)
    #             if key not in baidu_dict:
    #                 baidu_dict[key] = [score]
    #             else:
    #                 baidu_dict[key].append(score)
    #     except Exception as e:
    #         print(e)
    # df = pd.DataFrame(baidu_dict)
    # df.to_csv(os.path.join(root,"outputs","baidu_nrss.csv"),index=False)
    #
    # google_dict = {"images":[]}
    # for image in google_images:
    #     print(image)
    #     google_dict["images"].append(os.path.split(image)[-1])
    #     try:
    #         for key, val in funcs.items():
    #             score = val(image)
    #             if key not in google_dict:
    #                 google_dict[key] = [score]
    #             else:
    #                 google_dict[key].append(score)
    #     except Exception as e:
    #         print(e)
    # df = pd.DataFrame(google_dict)
    # df.to_csv(os.path.join(root, "outputs", "google_nrss.csv"), index=False)
    #
    # bing_dict = {"images":[]}
    # for image in bing_images:
    #     print(image)
    #     bing_dict["images"].append(os.path.split(image)[-1])
    #     try:
    #         for key, val in funcs.items():
    #             score = val(image)
    #             if key not in bing_dict:
    #                 bing_dict[key] = [score]
    #             else:
    #                 bing_dict[key].append(score)
    #     except Exception as e:
    #         print(e)
    # df = pd.DataFrame(bing_dict)
    # df.to_csv(os.path.join(root, "outputs", "bing_nrss.csv"), index=False)

if __name__=="__main__":
    main()